#include "PCFG_v2.h"
#include <chrono>
#include <fstream>
#include "SIMD_md5.h"
#include <iomanip>
#include <unordered_set>

#include <mpi.h>
using namespace std;
using namespace chrono;

// 编译指令如下
// g++ main.cpp train.cpp guessing.cpp md5.cpp -o main
// g++ main.cpp train.cpp guessing.cpp md5.cpp -o main -O1
// g++ main.cpp train.cpp guessing.cpp md5.cpp -o main -O2

// g++ correctness_guess.cpp train.cpp guessing.cpp md5.cpp -o correctness_guess
// g++ correctness_guess.cpp train.cpp guessing.cpp md5.cpp -o correctness_guess -O1
// g++ correctness_guess.cpp train.cpp guessing.cpp md5.cpp -o correctness_guess -O2

// g++ correctness_guess.cpp train.cpp guessing.cpp md5.cpp ./Pthread_OpenMP/Pthread_jin.cpp -o correctness_guess -pthread
// g++ correctness_guess.cpp train.cpp guessing.cpp md5.cpp  -o correctness_guess -fopenmp -O2

// mpicxx correctness_guess.cpp train.cpp guessing.cpp SIMD_md5.cpp -o correctness_guess  -fopenmp -O2

int main()
{
    // 打开文件以写入时间记录
    // ofstream time_records_file("./output/多线程/pthread_o1.txt");
    // if (!time_records_file.is_open()) {
    //     cerr << "无法打开文件 time_records.txt 进行写入！" << endl;
    //     return 1;
    // }

    double time_hash = 0;  // 用于MD5哈希的时间
    double time_guess = 0; // 哈希和猜测的总时长
    double time_train = 0; // 模型训练的总时长
    PriorityQueue q;
    auto start_train = system_clock::now();
    q.m.train("/guessdata/Rockyou-singleLined-full.txt");
    q.m.order();
    auto end_train = system_clock::now();
    auto duration_train = duration_cast<microseconds>(end_train - start_train);
    time_train = double(duration_train.count()) * microseconds::period::num / microseconds::period::den;


    
    // 加载一些测试数据
    unordered_set<std::string> test_set;
    ifstream test_data("/guessdata/Rockyou-singleLined-full.txt");
    int test_count=0;
    string pw;
    while(test_data>>pw)
    {   
        test_count+=1;
        test_set.insert(pw);
        if (test_count>=1000000)
        {
            break;
        }
    }
    int cracked=0;

    q.init();
    cout << "here" << endl;
    int curr_num = 0;
    auto start = system_clock::now();
    // 由于需要定期清空内存，我们在这里记录已生成的猜测总数
    int history = 0;

    // // 用于记录时间点和生成的猜测数量
    // vector<pair<int, double>> time_records;
    // // 记录上一次生成猜测的时间
    // auto last_time = start;

    // std::ofstream a("./files/results.txt");

    // 初始化 MPI 环境
    int argc = 0;
    char **argv = nullptr;
    MPI_Init(&argc, &argv);

    while (!q.priority.empty())
    {
        // // 记录开始时间
        // auto start_guess = system_clock::now(); 
        
        q.PopNext();

        // 记录结束时间
        // auto end_guess = system_clock::now();

        // 计算当前节点的猜测时间
        // auto duration_guess = duration_cast<microseconds>(end_guess - start);
        // double guess_time = double(duration_guess.count()) * microseconds::period::num / microseconds::period::den;
        // time_records.push_back({q.total_guesses + history, guess_time-time_hash});

        q.total_guesses = q.guesses.size();
        if (q.total_guesses - curr_num >= 100000)
        { 
            auto now = system_clock::now();
            auto duration = duration_cast<microseconds>(now - start);
            double elapsed_time = double(duration.count()) * microseconds::period::num / microseconds::period::den;

            cout << "Guesses generated: " <<history + q.total_guesses << endl;
            curr_num = q.total_guesses;

            // 在此处更改实验生成的猜测上限
            int generate_n=10000000;
            if (history + q.total_guesses > generate_n)
            {
                auto end = system_clock::now();
                auto duration = duration_cast<microseconds>(end - start);
                time_guess = double(duration.count()) * microseconds::period::num / microseconds::period::den;
                cout << "Guess time:" << time_guess - time_hash << "seconds"<< endl;
                cout << "Hash time:" << time_hash << "seconds"<<endl;
                cout << "Train time:" << time_train <<"seconds"<<endl;
                cout<<"Cracked:"<< cracked<<endl;

                break;
            }
        }
        // 为了避免内存超限，我们在q.guesses中口令达到一定数目时，将其中的所有口令取出并且进行哈希
        // 然后，q.guesses将会被清空。为了有效记录已经生成的口令总数，维护一个history变量来进行记录
        if (curr_num > 1000000)
        {
            auto start_hash = system_clock::now();

            
            // {
            //   alignas(16) bit32 state[4][4];
            //   string pw[4];
            //   // for (string pw : q.guesses)
            //   cout<<q.guesses.size()<<endl;

            //       for (int i = 0; i < q.guesses.size(); i += 4)
            //   {
            //     for (int j = i; j < q.guesses.size()&&(j-i)<4;j++){
            //       pw[j-i] = q.guesses[j];
            //       if (test_set.find(pw[j-i]) != test_set.end()) {
            //           cracked+=1;
            //       }
            //     }
            //     SIMD_MD5Hash(pw, state);
            //   }
            // }

            {
              alignas(16) bit32 state[4][4];
              string pw[4];

// 使用 OpenMP 并行化外层循环
#pragma omp parallel  for num_threads(7) shared(q, test_set) private(pw, state)
        for (int i = 0; i < q.guesses.size(); i += 4)
        {
          for (int j = i; j < q.guesses.size() && (j - i) < 4; j++)
          {
            pw[j - i] = q.guesses[j];
            if (test_set.find(pw[j - i]) != test_set.end())
            {
#pragma omp atomic
              cracked += 1; // 使用原子操作保证线程安全
            }
          }
          SIMD_MD5Hash(pw, state); // 假设 SIMD_MD5Hash 是线程安全的
        }
      }

      // 在这里对哈希所需的总时长进行计算
      auto end_hash = system_clock::now();
      auto duration = duration_cast<microseconds>(end_hash - start_hash);
      time_hash += double(duration.count()) * microseconds::period::num / microseconds::period::den;

      // 记录已经生成的口令总数
      history += curr_num;
      curr_num = 0;
      q.guesses.clear();
    }
  }

  // 清理 MPI 环境
  MPI_Finalize();
}
